﻿// main.cpp
// 本项目的主程序入口。
// 负责协调整个优化流程，包括：
// 1. 设置算法参数。
// 2. 加载和预处理数据。
// 3. 生成初始调度方案。
// 4. 执行两阶段混合优化算法。
// 5. 总结并保存最终结果。

#include <iostream>
#include <vector>
#include <string>
#include <fstream>
#include <iomanip>
#include <omp.h>
#include <thread>
#include <random>
#include <direct.h>  // Windows目录创建
#include <sys/stat.h>

#include "DataTypes.h"
#include "Graph.h"
#include "Scoring.h"
#include "Validation.h"
#include "InitialSolution.h"
#include "OptimizerPhase1.h"
#include "OptimizerPhase2.h"

// 创建结果目录的辅助函数 (Windows兼容)
void create_result_directory(const std::string& case_name) {
    std::string result_dir = "result";
    std::string case_dir = result_dir + "/" + case_name;
    
    // 创建主result目录
    struct stat st = {0};
    if (stat(result_dir.c_str(), &st) == -1) {
        _mkdir(result_dir.c_str());
    }
    
    // 创建案例子目录
    if (stat(case_dir.c_str(), &st) == -1) {
        _mkdir(case_dir.c_str());
    }
    
    std::cout << "创建结果目录: " << case_dir << std::endl;
}

// 辅助函数：将一个长整型向量写入到CSV文件中，用于后续可视化。
void write_vector_to_file(const std::string& filename, const std::vector<long long>& data) {
    std::ofstream outfile(filename);
    if (!outfile.is_open()) {
        std::cerr << "错误: 无法打开文件写入: " << filename << std::endl;
        return;
    }
    for (const auto& val : data) {
        outfile << val << "\n";
    }
    std::cout << "数据已写入 " << filename << std::endl;
}

// 辅助函数：将最终的调度方案（节点ID列表）写入到文本文件中。
void write_schedule_to_file(const std::string& filename, const std::vector<int>& schedule) {
    std::ofstream outfile(filename);
    if (!outfile.is_open()) {
        std::cerr << "错误: 无法打开文件写入: " << filename << std::endl;
        return;
    }
    for (const auto& val : schedule) {
        outfile << (val - 1) << "\n"; // 从1-based索引转换为0-based以便于其他工具使用
    }
}


int main() {
    // ========================================================================
    // 核心环境设置区
    // ========================================================================

    // --- 线程管理 ---
    // 获取硬件线程总数，并设置程序使用线程数为 (总数 - 1)，以保留一个核心给系统
    int max_threads = std::thread::hardware_concurrency();
    int desired_threads = max_threads > 1 ? max_threads - 1 : 1; // 确保至少有1个线程
    omp_set_num_threads(desired_threads);

    // --- 随机数管理 ---
    // 创建一个唯一的、使用固定种子的随机数生成器，以确保结果可复现
    unsigned int fixed_seed = 12345;
    std::mt19937 gen(fixed_seed);

    // --- 诊断信息打印 ---
#ifdef _OPENMP
    fprintf(stdout, "OpenMP 已启用。\n");
    fprintf(stdout, "硬件线程总数: %d, 本次运行使用线程数: %d\n", max_threads, desired_threads);
#else
    fprintf(stderr, "警告: OpenMP 未启用。程序将以单线程模式运行。\n");
    fprintf(stderr, "      请检查您的编译器标志 (例如，在VS中为 /openmp，在GCC/Clang中为 -fopenmp)。\n");
#endif
    fprintf(stdout, "========================================================================\n\n");

    // ========================================================================
    // 设置区: 在这里调整算法参数
    // ========================================================================

    // 测试案例列表 (先用单个案例测试)
    std::vector<std::string> test_cases = {
        //"Conv_Case0",
        "Conv_Case1",
        "Matmul_Case0",
        "Matmul_Case1",
        "FlashAttention_Case1"
    };
    
    // 处理所有案例
    for (const auto& case_name : test_cases) {
        std::cout << "\n\n========================================" << std::endl;
        std::cout << "开始处理案例: " << case_name << std::endl;
        std::cout << "========================================" << std::endl;
        
        // 定义数据文件的路径
        const std::string nodes_filename = case_name + "_Nodes.csv";
        const std::string edges_filename = case_name + "_Edges.csv";

        // 从输入文件名中提取基础名称，用于生成输出文件名
        std::string base_name = case_name;
        
        // 创建对应的结果目录
        create_result_directory(base_name);
        
        // 设置结果文件路径前缀
        std::string result_prefix = "result/" + base_name + "/" + base_name;

        // --- 阶段一: 广域探索参数 (随机移动策略) ---
        // 阶段一的目标是利用并行化和随机性，快速地对整个解空间进行广泛搜索，
        // 以找到一个比较好的"有潜力的区域"，为第二阶段的精细搜索打下基础。
        Options1 options1;

    // 迭代总次数。这是阶段一搜索量的主要控制参数。
    // 增加此值可以进行更广泛的搜索，可能找到更好的初始解，但会增加运行时间。
    options1.iterations = 50000;

    // 提前终止的"耐心值"。如果在连续这么多次迭代中，分数都没有任何改善，
    // 算法会提前结束阶段一，避免浪费时间。
    options1.patience = 5000;

    // 并行生成的候选解数量（批处理大小）。
    // 程序会利用OpenMP进行并行计算，此参数定义了每一轮迭代中要并行评估多少个不同的移动方案。
    // 通常设置为CPU核心数的倍数以最大化效率。
    options1.batch_size = desired_threads * 10;


    // --- 阶段二: 迭代局部搜索(ILS)参数 ---
    // 阶段二的核心思想是从阶段一找到的解出发，进行一次深入的、系统性的局部搜索以找到“局部最优解”。
    // 然后，通过“扰动”（Perturbation）打乱当前最优解，跳出局部最优，再进行新一轮的局部搜索。
    // 如此重复，以期找到更好的局部最优，并逐步逼近全局最优解。
    Options2 options2;

    // 局部搜索模块中的最大轮次(Pass)。在一个局部搜索中，算法会完整地遍历一遍所有可优化的节点。
    // 如果一轮遍历后找到了更优解，它会开始新的一轮(Pass)。此参数限制了连续改进的最大轮次数，
    // 防止在一个局部最优的“斜坡”上花费过多时间。
    options2.max_passes = 500;

    // 整个阶段二的总运行时长上限（分钟）。这是一个硬性的时间限制，确保算法能在规定时间内完成。
    options2.time_limit_minutes = 30;

    // 【外层批处理】在局部搜索的每一轮中，并行处理的“焦点节点”批次大小的乘数。
    // 例如，如果有16个核心，此值为1，则一次并行处理16个不同的焦点节点。
    options2.num_focus_nodes_multiplier = 2;

    // 【内层批处理】对于每一个“焦点节点”，它可能有很多个可以移动到的新位置。
    // 此参数控制一次并行评估多少个这样的“候选移动方案”。
    options2.batch_size_multiplier = 100;

    // ILS的“扰动-搜索”循环次数。这是ILS框架的核心。
    // 0表示只进行一次从初始解开始的局部搜索。大于0的值表示在找到第一个局部最优解后，
    // 再进行N次“扰动当前最优解 -> 开始新一轮局部搜索”的循环。
    // 这是算法跳出局部最优、寻找全局最优的关键。增加此值能提升找到更优解的概率，但会显著增加运行时间。
    options2.num_perturbations = 1000;


    // ========================================================================
    // 第一、二部分：加载数据与预处理
    // ========================================================================
    std::cout << "--- Section 1&2: Loading and pre-processing data ---" << std::endl;
    Graph graph;
    CodeMap code_map;
    try {
        // 从CSV文件读取节点和边的信息，构建图的数据结构
        if (!graph.load_from_csv(nodes_filename, edges_filename, code_map)) { return 1; }
    }
    catch (const std::exception& e) {
        std::cerr << "错误：加载CSV文件时发生异常: " << e.what() << std::endl;
        return 1;
    }
    std::cout << "数据加载与预处理完成。" << std::endl;


    // ========================================================================
    // 第三部分：生成初始可行解
    // ========================================================================
    std::cout << "--- Section 3: Generating initial solution ---" << std::endl;
    std::vector<int> initial_schedule;
    try {
        // 调用函数生成一个满足所有硬性约束（依赖约束和L0约束）的初始调度方案
        initial_schedule = InitialSolution::generate_l0_constrained_schedule(graph, code_map);
    }
    catch (const std::exception& e) {
        std::cerr << "错误: " << e.what() << std::endl;
        return 1;
    }
    // 计算这个初始方案的得分（峰值内存）和内存使用历史
    auto initial_res = Scoring::calculate_max_vstay(initial_schedule, graph, code_map);
    long long initial_score = initial_res.first;
    auto initial_history = initial_res.second;
    fprintf(stdout, "初始可行解的峰值内存占用为: %lld\n", initial_score);


    // ========================================================================
    // 第四部分：执行两阶段混合优化
    // ========================================================================
    std::cout << "--- Section 4: Executing the two-phase hybrid optimization ---" << std::endl;

    // --- 阶段一：执行广域探索 ---
    fprintf(stdout, "\n--- 开始阶段一: 广域探索 ---\n");
    auto phase1_res = OptimizerPhase1::optimizer_random_sampling_bottleneck1(
        initial_schedule, initial_history, initial_score, graph, code_map, options1, gen);
    auto schedule_phase1 = phase1_res.schedule;
    long long score_phase1 = phase1_res.best_score;
    auto phase1_convergence = phase1_res.convergence_history;
    auto phase1_iterations = phase1_res.improvement_iterations;
    fprintf(stdout, "--- 阶段一完成! 分数从 %lld 优化至 %lld (共 %zu 次改进) ---\n", 
            initial_score, score_phase1, phase1_convergence.size());

    // 对阶段一的输出进行验证，确保它没有产生非法解
    fprintf(stdout, "--- 验证阶段一输出结果... ---\n");
    if (Validation::is_schedule_fully_valid(schedule_phase1, graph, code_map)) {
        fprintf(stdout, "阶段一输出方案满足所有约束。\n");
    }
    else {
        fprintf(stderr, "致命错误：阶段一优化器产出了不满足约束的解！程序终止。\n");
        return 1;
    }

    // --- 阶段二：执行迭代局部搜索 (ILS) ---
    fprintf(stdout, "\n--- 开始阶段二: 迭代局部搜索 (ILS) ---\n");
    auto history_phase1 = Scoring::calculate_max_vstay(schedule_phase1, graph, code_map).second;
    auto final_res = OptimizerPhase2::optimizer_systematic_search_bottleneck(
        schedule_phase1, history_phase1, score_phase1, graph, code_map, options2, gen);
    auto final_schedule = final_res.schedule;
    long long final_score = final_res.best_score;
    auto phase2_convergence = final_res.convergence_history;
    auto phase2_iterations = final_res.improvement_iterations;
    fprintf(stdout, "--- 阶段二(ILS)完成! 分数从 %lld 进一步优化至 %lld (共 %zu 次改进) ---\n", 
            score_phase1, final_score, phase2_convergence.size());

    // 对最终结果进行验证，确保它没有产生非法解
    fprintf(stdout, "--- 验证阶段二输出结果... ---\n");
    if (Validation::is_schedule_fully_valid(final_schedule, graph, code_map)) {
        fprintf(stdout, "阶段二输出方案满足所有约束。\n");
    }
    else {
        fprintf(stderr, "致命错误：阶段二优化器产出了不满足约束的解！程序终止。\n");
        return 1;
    }


    // ========================================================================
    // 第五、六、七部分：结果总结、数据导出与保存
    // ========================================================================
    std::cout << "--- Section 5,6,7: Final results, data export and saving ---" << std::endl;

    // 重新计算中间解和最终解的内存历史，用于后续的统计和数据导出
    auto history_phase1_final = Scoring::calculate_max_vstay(schedule_phase1, graph, code_map).second;
    auto final_history = Scoring::calculate_max_vstay(final_schedule, graph, code_map).second;

    // 在控制台打印最终的优化摘要
    fprintf(stdout, "\n\n--- 最终优化结果总结 ---\n");
    fprintf(stdout, "初始解分数: %lld\n", initial_score);
    fprintf(stdout, "阶段一后分数: %lld\n", score_phase1);
    fprintf(stdout, "最终优化分数: %lld\n", final_score);
    if (initial_score > 0) {
        fprintf(stdout, "总计内存占用降低了: %lld (%.2f%%)\n", initial_score - final_score,
            100.0 * (initial_score - final_score) / (double)initial_score);
    }

        // 将用于绘图的数据导出到CSV文件，方便使用Python, MATLAB, Excel等外部工具进行可视化
        write_vector_to_file(result_prefix + "_initial_history.csv", initial_history);
        write_vector_to_file(result_prefix + "_phase1_history.csv", history_phase1_final);
        write_vector_to_file(result_prefix + "_final_history.csv", final_history);
        
        // 导出收敛曲线数据 (新增功能)
        std::cout << "保存收敛曲线数据..." << std::endl;
        
        // 阶段1收敛历史
        std::ofstream phase1_conv_file(result_prefix + "_phase1_convergence.csv");
        phase1_conv_file << "iteration,score\n";
        for (size_t i = 0; i < phase1_convergence.size(); ++i) {
            phase1_conv_file << phase1_iterations[i] << "," << phase1_convergence[i] << "\n";
        }
        phase1_conv_file.close();
        std::cout << "阶段1收敛数据已写入 " << result_prefix + "_phase1_convergence.csv" << std::endl;
        
        // 阶段2收敛历史
        std::ofstream phase2_conv_file(result_prefix + "_phase2_convergence.csv");
        phase2_conv_file << "iteration,score\n";
        for (size_t i = 0; i < phase2_convergence.size(); ++i) {
            phase2_conv_file << phase2_iterations[i] << "," << phase2_convergence[i] << "\n";
        }
        phase2_conv_file.close();
        std::cout << "阶段2收敛数据已写入 " << result_prefix + "_phase2_convergence.csv" << std::endl;
        
        // 合并的整体收敛历史
        std::ofstream overall_conv_file(result_prefix + "_overall_convergence.csv");
        overall_conv_file << "phase,iteration,score\n";
        for (size_t i = 0; i < phase1_convergence.size(); ++i) {
            overall_conv_file << "1," << phase1_iterations[i] << "," << phase1_convergence[i] << "\n";
        }
        for (size_t i = 0; i < phase2_convergence.size(); ++i) {
            overall_conv_file << "2," << phase2_iterations[i] << "," << phase2_convergence[i] << "\n";
        }
        overall_conv_file.close();
        std::cout << "整体收敛数据已写入 " << result_prefix + "_overall_convergence.csv" << std::endl;

    int num_nodes = graph.get_num_nodes();
    std::vector<long long> initial_pos(num_nodes + 1, 0), final_pos(num_nodes + 1, 0);
    for (size_t i = 0; i < initial_schedule.size(); ++i) initial_pos[initial_schedule[i]] = i + 1;
    for (size_t i = 0; i < final_schedule.size(); ++i) final_pos[final_schedule[i]] = i + 1;

        std::vector<long long> displacement;
        std::ofstream pos_file(result_prefix + "_node_positions.csv");
        pos_file << "node_id,initial_pos,final_pos\n";
        for (int i = 1; i <= num_nodes; ++i) {
            pos_file << i << "," << initial_pos[i] << "," << final_pos[i] << "\n";
            displacement.push_back(final_pos[i] - initial_pos[i]);
        }
        pos_file.close();
        std::cout << "数据已写入 " << result_prefix + "_node_positions.csv" << std::endl;
        write_vector_to_file(result_prefix + "_displacement.csv", displacement);

        // 将最终优化好的调度方案保存到文本文件
        try {
            write_schedule_to_file(result_prefix + "_final_schedule_optimized.txt", final_schedule);
            fprintf(stdout, "最终优化方案已成功保存至: %s\n", (result_prefix + "_final_schedule_optimized.txt").c_str());
        }
        catch (const std::exception& e) {
            fprintf(stderr, "错误：无法保存最终方案到文件。%s\n", e.what());
        }

        // 生成案例汇总报告
        std::string summary_file = result_prefix + "_summary.txt";
        std::ofstream summary(summary_file);
        if (summary.is_open()) {
            summary << "=== GMCM 2025 Problem 1 Results Summary ===" << std::endl;
            summary << "Task Name: " << base_name << std::endl;
            summary << "Graph Size: " << graph.get_num_nodes() << " nodes" << std::endl;
            summary << "Initial V_stay: " << initial_score << std::endl;
            summary << "Phase 1 V_stay: " << score_phase1 << std::endl;
            summary << "Final V_stay: " << final_score << std::endl;
            if (initial_score > 0) {
                summary << "Improvement: " << (initial_score - final_score) 
                       << " (" << std::fixed << std::setprecision(2) 
                       << 100.0 * (initial_score - final_score) / (double)initial_score 
                       << "%)" << std::endl;
            }
            summary << "Generated Files:" << std::endl;
            summary << "  - " << base_name << "_final_schedule_optimized.txt (调度序列)" << std::endl;
            summary << "  - " << base_name << "_initial_history.csv (初始内存历史)" << std::endl;
            summary << "  - " << base_name << "_phase1_history.csv (阶段1内存历史)" << std::endl;
            summary << "  - " << base_name << "_final_history.csv (最终内存历史)" << std::endl;
            summary << "  - " << base_name << "_node_positions.csv (节点位置变化)" << std::endl;
            summary << "  - " << base_name << "_displacement.csv (节点位移数据)" << std::endl;
            summary.close();
            std::cout << "汇总报告已生成: " << summary_file << std::endl;
        }

        std::cout << "案例 " << case_name << " 处理完成" << std::endl;
    }  // 结束案例循环
    
    std::cout << "\n========================================" << std::endl;
    std::cout << "--- 所有案例处理完成 ---" << std::endl;
    std::cout << "========================================" << std::endl;

    return 0;
}